# llms/ollama_llm.py
import ollama
from typing import Optional
from config import app_config
from .llm_base import BaseLLM
from .llm_registry import register_llm
from chat_memory import build_message_history, append_to_history
from ..repository.characterCard_repository import character_card_repo
from ..repository.chatSession_repository import chat_session_repo


@register_llm("ollama")
class OllamaLLM(BaseLLM):
    def __init__(self) -> None:
        super().__init__()
        # 根据配置初始化 Ollama 客户端
        self._client = ollama.Client(host=app_config.OLLAMA_HOST)
        self._model_name = app_config.OLLAMA_MODEL
        self._history_limit = 5

    """Ollama LLM 实现，支持记忆、流式输出、会话管理"""

    def generate(
        self,
        prompt: str,
        session_id: int = -1,
        max_tokens: int = 2048,
        stream: bool = False,
        enable_memory: bool = False,
        **kwargs  # 允许扩展其他参数
    ) -> Optional[str]:
        """
        调用本地 Ollama 模型，支持记忆、流式输出。
        Args:
            prompt: 用户输入
            session_id: 会话 ID，用于区分不同用户或对话
            system_prompt: 系统提示词
            temperature: 生成温度
            max_tokens: 最大生成 token 数
            stream: 是否流式输出
            enable_memory: 是否启用记忆（上下文）

        Returns:
            AI 回复文本，失败返回 None
        """

        try:
            # ===== 0. 预设参数 =====
            session_data = chat_session_repo.get_session_by_id(session_id)
            if session_data and session_data.get("character_id"):
                character_id = session_data["character_id"]
                # 根据角色卡生成系统提示词
                system_prompt = character_card_repo.build_system_prompt(character_id)
                # 根据会话读取角色卡中的指定的模型参数
                model_params = character_card_repo.get_model_params(session_id)
                if not model_params and model_params["temperature"] and model_params["top_p"]:
                    temperature = model_params["temperature"]
                    top_p = model_params["top_p"]
                else:
                    raise ValueError(f"会话 {session_id} 关联的角色卡缺少模型参数")
            else:
                raise ValueError(f"会话 {session_id} 未关联角色卡")

            # ===== 1. 构建消息历史 =====
            messages = build_message_history(
                session_id=session_id,
                system_prompt=system_prompt,
                user_input=prompt,
                enable_memory=enable_memory,
                history_limit=self._history_limit  # 根据实例属性设置历史长度
            )

            # ===== 2. 调用 Ollama =====
            response = self._client.chat(
                model=self._model_name,
                messages=messages,
                stream=stream,
                options={
                    "temperature": temperature,
                    "num_predict": max_tokens,
                    "top_p": top_p,
                },
            )

            # ===== 3. 处理响应（带成功判断）=====
            reply = ""
            if not stream:
                # 非流式响应判断
                if  "error" in response: # type: ignore
                    raise ValueError(f"Ollama 请求失败：{response['error']}") # type: ignore
                if not all(key in response for key in ["message", "done"]) or not response["done"]: # type: ignore
                    raise RuntimeError(f"Ollama 响应异常：{response}")
                reply = response["message"]["content"] # type: ignore
            else:
                # 流式响应判断（依赖 _handle_streaming_response 内部判断）
                reply = self._handle_streaming_response(response)
            
            return reply

        except Exception as e:
            print(f"🚨 Ollama 调用异常: {str(e)}")
            return "抱歉，当前无法生成回复，请检查 Ollama 服务是否正常~"

    def cleanup(self):
        """Ollama 无需显式释放资源"""
        pass

    # ==============================
    # 私有辅助函数
    # ==============================
    def _handle_streaming_response(self,stream_response) -> str:
        """处理流式响应，附带成功判断"""
        full_content = []
        has_error = False
        error_msg = ""
        
        for chunk in stream_response:
            # 1. 检查当前片段是否有错误
            if "error" in chunk:
                has_error = True
                error_msg = chunk["error"]
                break  # 中断流处理
            # 2. 提取当前片段的内容（流式响应中，content 可能分多次返回）
            if "message" in chunk and "content" in chunk["message"]:
                full_content.append(chunk["message"]["content"])
            # 3. 遇到 done=True 时终止循环（正常结束）
            if chunk.get("done", False):
                break
        
        # 4. 流处理结束后，判断是否成功
        if has_error:
            raise ValueError(f"Ollama 流式请求失败：{error_msg}")
        if not full_content:  # 无错误但无内容（如流中断）
            raise RuntimeError("Ollama 流式响应为空，可能服务中断")
        
        return "".join(full_content)

if __name__ == "__main__":
    pass
    # 测试代码
    # llm = OllamaLLM()
    # response = llm.generate("你好，介绍一下你自己吧！", session_id=1, stream=True, enable_memory=True)
    # print("最终回复：", response)